#include "PLUSolver.h"
#include <math.h>

template<typename T>
PLUSolver<T>::PLUSolver() {}

template<typename T>
PLUSolver<T>::~PLUSolver() {
    if(Y!=NULL) free(Y);
    if(rL!=NULL) free(rL);
    if(rU!=NULL) free(rU);
    if(L!=NULL) free(L);
    if(U!=NULL) free(U);
}

template<typename T>
PLUSolver<T>::PLUSolver(T** _A, int _n, T* _b, int* _P) {
    A = _A;
    n = _n;
    b = _b;
    P = _P;
    Y = (T*)malloc(sizeof(T)*n);
    for (int i = 0; i < n; i++){
        P[i] = i;
        Y[i] = 0;
    }
    rL = NULL;
    rU = NULL;
    L = NULL;
    U = NULL;
}

template<typename T>
int PLUSolver<T>::PLU(){

    for(int i = 0; i < n; i++){
        T pivot = 0;
        int posPivot;
        for (int j = i; j < n; j++){
            if (abs(pivot) < abs(A[j][i])){
                pivot = A[j][i];
                posPivot = j;
            }
        }

        swap(i, posPivot);
        for (int j = i+1; j < n; j++){
            // A[j][i]
            A[j][i] = A[j][i]/A[i][i];
            for (int k = i + 1; k < n; k++) {
                A[j][k] -=  A[j][i] * A[i][k];
            }
        }

        // PrintA();
    }
    return 1;
}

template<typename T>
int PLUSolver<T>::Solve(T* _X){
    // Ly = b forward
    // Y[0] = b[0] / (1.0);
    // Y[1] = b[1] - (Y[0] * A[1][0])
    // 
    X = _X;
    for(int i = 0; i < n; i++){
        T leftSum = 0;
        for(int j = 0; j < i; j++){
            leftSum += Y[j] * A[i][j];
        }
        Y[i] = (b[i] - leftSum);
    }

    // Ux = y backward
    // X[n-1] = Y[n-1] / A[n-1][n-1]
    // X[n-2] = (Y[n-2] - (X[n-1] * A[n-2][n-1])) / A[n-2][n-2]
    // X[n-3] = (Y[n-3] - (X[n-1] * A[n-3][n-1] + X[n-2] * A[n-3][n-2])) / A[n-3][n-3]
    for(int i = (n-1); i >= 0; i--){
        T left_sum = 0;
        for(int j = (n-1); j > i; j--){
            left_sum += X[j] * A[i][j];
        }
        X[i] = (Y[i] - left_sum) / A[i][i];
    }

    return 1;
}

template<typename T>
int PLUSolver<T>::reverseLU(){
    rL = (T**)malloc(sizeof(T*) * n);
    L = (T**)malloc(sizeof(T*) * n);
    rU = (T**)malloc(sizeof(T*) * n);
    U = (T**)malloc(sizeof(T*) * n);
    for(int i = 0; i < n; i++){
        rL[i] = (T*)malloc(sizeof(T) * n);
        L[i] = (T*)malloc(sizeof(T) * n);
        rU[i] = (T*)malloc(sizeof(T) * n);
        U[i] = (T*)malloc(sizeof(T) * n);
    }

    // cal rL
    // rL * L = I
    for(int i = 0; i < n; i++){
        for(int j = (n-1); j >= 0; j--){
            if(j > i){
                rL[i][j] = 0;
            }
            else if(j == i){
                rL[i][i] = (T)((1.0)/(1.0));
            }
            else{
                T upSum = 0;
                for (int k = (j + 1); k <= i; k++){
                    upSum += rL[i][k] * A[k][j]; // k < j
                }
                upSum = -upSum;
                rL[i][j] = upSum;
            }   
        }
    }

    // cal rU
    // rU * U = I
    for(int i = (n-1); i >= 0; i--){
        for(int j = 0; j < n; j++){
            if(j < i){
                rU[i][j] = 0;
            }
            else if(j == i){
                rU[i][j] = (1.0) / A[i][j];
            }
            else{ // j > i
                T upSum = 0;
                for(int k = i; k < j; k++){
                    upSum += rU[i][k] * A[k][j];
                }
                rU[i][j] = - upSum / A[j][j];
            }
        }
    }

    return 1;
}

template<typename T>
int PLUSolver<T>::GetReverseA(T** _rA){
    rA = _rA;
    reverseLU();

    T** temp = (T**)malloc(sizeof(T*) * n);
    for (int i = 0; i < n; i++){
        temp[i] = (T*)malloc(sizeof(double) * n);
    }

    // 按照题目要求，rA是我们要求的东西，所以应该在外部申明
    gemm(rU, rL, temp, n);

    T** matrixP = (T**)malloc(sizeof(T*) * n);
    for (int i = 0; i < n; i++){
        matrixP[i] = (T*)malloc(sizeof(double) * n);
    }

    for (int i = 0; i < n; i++){
        int pos = P[i];
        for (int j = 0; j < n; j++){
            matrixP[i][j] = 0;
        }
        matrixP[i][pos] = 1.0;
    }

    gemm(temp, matrixP, rA, n);

    // delete
    for(int i = 0; i < n; i++){
        free(temp[i]);
    }
    free(temp);

    return 1;
}

template<typename T>
int PLUSolver<T>::gemm(T** M1, T** M2, T** res, int n){
    if(res == NULL || M1 == NULL || M2 == NULL) return 0;
    for (int i = 0; i < n; i++){
        for (int j = 0; j < n; j++){
            res[i][j] = 0;
            for (int k = 0; k < n; k++){
                res[i][j] += M1[i][k] * M2[k][j];
            }
        }
    }
    return 1; 
}


template<typename T>
int PLUSolver<T>::swap(int row1, int row2){
    T* temp;
    temp = A[row1];
    A[row1] = A[row2];
    A[row2] = temp;
    int tempP;
    tempP = P[row1];
    P[row1] = P[row2];
    P[row2] = tempP;
    T tempB;
    tempB = b[row1];
    b[row1] = b[row2];
    b[row2] = tempB;
    return 1;
}

template<typename T>
int PLUSolver<T>::PrintMatrix(T** M, int size, std::string name){
    std::cout<<"Print "<<name<<" ----------------------------------------"<<std::endl;
    for (int i = 0; i < n; i++){
        for (int j = 0; j < n; j++){
            printf("%lf ", M[i][j]);
        }
        printf(" | %d \n", P[i]);
    }
    std::cout<<"-----------------------------------------------------------"<<std::endl<<std::endl;
    return 1;
}

template<typename T>
int PLUSolver<T>::PrintVector(T* Vec, int size, std::string name){
    std::cout<<"Print "<<name<<" ----------------------------------------"<<std::endl;
    
    for (int i = 0; i < size; i++){
        printf("%lf \n", Vec[i]);
    }

    std::cout<<"-----------------------------------------------------------"<<std::endl<<std::endl;
    return 1;
}

